# Author : ZZH
# Date : 2025/5/23
import argparse
import os
import parser
import platform
import time
import logging
from datetime import datetime
from pathlib import Path
from collections import defaultdict
from tqdm import tqdm
import cv2
import torch
# from utils.utils import rename_log_file
from ultralytics import YOLO


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='YOLOv8/YOLOv11推理脚本')
    parser.add_argument('--source', type=str, default=r'../demo2.mp4', help='输入源（图像/文件夹/视频）')
    parser.add_argument('--weights', type=str, default=r'../models/checkpoints/train5-20250523-131504-yolov8m-best.pt',
                        help='模型权重路径')
    parser.add_argument('--imgsz', type=int, default=640, help='推理尺寸')
    parser.add_argument('--device', type=str, default='cpu', help='计算设备 (e.g. 0/cpu)')
    parser.add_argument('--base-path', type=str, default='../', help='基础存储路径')
    # 添加视频处理参数
    parser.add_argument('--video', type=str,
                        help='Path to input video for inference')
    parser.add_argument('--show', action='store_true',
                        help='实时显示处理结果')

    return parser.parse_args()


def init_logger(base_path, temp_name):
    """初始化日志记录器"""
    log_dir = os.path.join(base_path, 'logging', 'infer')
    os.makedirs(log_dir, exist_ok=True)
    log_path = os.path.join(log_dir, temp_name)

    logger = logging.getLogger('YOLO_Inference')
    logger.setLevel(logging.INFO)

    # 清理现有处理器
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)

    # 文件处理器
    file_handler = logging.FileHandler(log_path)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)

    return logger, log_path


def create_save_dir(base_path):
    """创建递增结果目录"""
    base_dir = os.path.join(base_path, 'runs', 'infer')
    os.makedirs(base_dir, exist_ok=True)

    # 获取现有目录的最大编号
    existing = [d for d in os.listdir(base_dir) if d.startswith('predict')]
    n = 1
    if existing:
        nums = [int(d[7:]) for d in existing if d[7:].isdigit()]
        if nums: n = max(nums) + 1

    save_dir = os.path.join(base_dir, f'predict{n}')
    os.makedirs(save_dir, exist_ok=True)
    return save_dir, n


def get_sample_count(source):
    """获取输入样本数量"""
    path = Path(source)
    if not path.exists():
        raise FileNotFoundError(f"路径不存在: {source}")

    if path.is_file():
        return 1
    elif path.is_dir():
        exts = ['.jpg', '.jpeg', '.png', '.bmp', '.mp4', '.avi']
        return len([f for f in path.iterdir() if f.suffix.lower() in exts])
    return 0


def process_video_with_realtime_display(model, video_path, output_dir, conf_thresh=0.5, show=False):
    """
    支持实时显示的视频推理处理
    :param model: 加载的YOLO模型
    :param video_path: 输入视频路径
    :param output_dir: 输出目录
    :param conf_thresh: 置信度阈值
    :param show: 是否实时显示画面
    :return: 处理后视频路径
    """
    # 创建输出目录
    output_dir.mkdir(parents=True, exist_ok=True)

    # 视频输入设置
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise IOError(f"无法打开视频文件 {video_path}")

    # 获取视频参数
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 视频输出设置
    output_path = output_dir / f"vis_{video_path.name}"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(
        str(output_path),
        fourcc,
        fps,
        (frame_width, frame_height)
    )

    # 进度条设置
    pbar = tqdm(total=total_frames, desc='Processing Video', unit='frame')

    # 实时显示窗口设置
    if show:
        cv2.namedWindow('YOLOv8 Real-time Detection', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('YOLOv8 Real-time Detection', frame_width // 2, frame_height // 2)

    # 性能计数器
    fps_counter = {
        'frames': 0,
        'start': time.time(),
        'fps': 0
    }

    # 逐帧处理
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # 执行推理
        start_infer = time.time()
        results = model(frame, conf=conf_thresh, verbose=False)
        infer_time = time.time() - start_infer

        # 可视化结果
        annotated_frame = results[0].plot(
            line_width=2,
            font_size=0.8,
            labels=True,
            conf=True
        )

        # 添加性能信息
        fps_counter['frames'] += 1
        if (time.time() - fps_counter['start']) > 1:
            fps_counter['fps'] = fps_counter['frames'] / (time.time() - fps_counter['start'])
            fps_counter['frames'] = 0
            fps_counter['start'] = time.time()

        cv2.putText(annotated_frame,
                    f"FPS: {fps_counter['fps']:.1f} | Infer: {infer_time * 1000:.1f}ms",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (0, 255, 0), 2)

        # 写入输出视频
        out.write(annotated_frame)

        # 实时显示
        if show:
            cv2.imshow('YOLOv8 Real-time Detection', annotated_frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        pbar.update(1)

    # 释放资源
    cap.release()
    out.release()
    pbar.close()
    if show:
        cv2.destroyAllWindows()

    return output_path


def main():
    args = parse_args()
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

    # 初始化临时日志
    logger, log_path = init_logger(args.base_path, f'temp-{timestamp}-yolov8n-best.log')
    logger.info("推理开始")

    try:
        # 文件存在性检查
        if not os.path.exists(args.source):
            raise FileNotFoundError(f"输入源不存在: {args.source}")
        if not os.path.exists(args.weights):
            raise FileNotFoundError(f"权重文件不存在: {args.weights}")

        # 创建结果目录
        save_dir, run_id = create_save_dir(args.base_path)

        # 记录设备信息
        device_info = "CPU"
        if args.device.lower() != 'cpu' and torch.cuda.is_available():
            dev_id = int(args.device) if args.device.isdigit() else 0
            dev_props = torch.cuda.get_device_properties(dev_id)
            device_info = f"GPU={dev_props.name} ({dev_props.total_memory / 1024 ** 3:.2f}GB)"

        logger.info(f"设备信息: OS={platform.platform()}, {device_info}")
        logger.info(f"输入信息: 路径={args.source}, 样本数={get_sample_count(args.source)}")
        logger.info(f"模型信息: {os.path.basename(args.weights)}")

        # 加载模型
        model = YOLO(args.weights)
        device = f'cuda:{args.device}' if args.device != 'cpu' and torch.cuda.is_available() else 'cpu'

        # 执行推理
        start = time.time()
        results = model.predict(
            source=args.source,
            imgsz=args.imgsz,
            device=device,
            save=True,
            save_dir=save_dir,
            exist_ok=True,
            stream=True
        )
        results = list(results)  # 转换为列表以便统计
        total_time = time.time() - start

        # 统计检测结果
        total_boxes = 0
        class_counts = defaultdict(int)
        for r in results:
            if r.boxes:
                total_boxes += len(r.boxes)
                for cls in r.boxes.cls:
                    class_name = model.names[int(cls)]
                    class_counts[class_name] += 1

        # 记录结果
        logger.info(f"推理结果: 检测框总数={total_boxes}, " +
                    ", ".join([f"{k}={v}" for k, v in class_counts.items()]))
        logger.info(f"结果目录: {save_dir}")
        logger.info(f"耗时: 总时间={total_time:.2f}秒, 平均每项={total_time / max(1, len(results)):.3f}秒")

        # 控制台输出
        print("\n推理结果：")
        print(f"- 检测框总数: {total_boxes}")
        for cls, count in class_counts.items():
            print(f"- {cls}: {count}")

        # # 重命名日志文件
        # new_log_name = f'infer{run_id}-{timestamp}-yolov8n-best.log'
        # os.rename(log_path, os.path.join(os.path.dirname(log_path), new_log_name))
        # logger.info(f"日志重命名为: {new_log_name}")

    except Exception as e:
        logger.error(f"推理失败: {str(e)}", exc_info=True)
        raise
    finally:
        logging.shutdown()

    if args.video:
        video_path = Path(args.video)
        if not video_path.exists():
            raise FileNotFoundError(f"视频文件不存在: {video_path}")

        # 处理视频并显示
        output_video = process_video_with_realtime_display(
            model=model,
            video_path=video_path,
            output_dir='../out_zhen',
            conf_thresh=0.5,
            show=args.show
        )

        logger.info(f"视频处理完成: {output_video}")


if __name__ == '__main__':
    main()
